package compiler;

import grammar.ApplicationTree;
import grammar.FunLexer;
import grammar.FunParser;
import org.antlr.v4.runtime.ParserRuleContext;
import org.antlr.v4.runtime.tree.ParseTree;
import org.antlr.v4.runtime.tree.TerminalNode;
import org.objectweb.asm.*;
import sun.reflect.generics.reflectiveObjects.NotImplementedException;

import java.util.Map;
import java.util.Set;
import java.util.TreeMap;

import static org.objectweb.asm.Opcodes.*;

public class LetCompiler extends TestCompileVisitor {


    public LetCompiler(FunParser.LetContext context, Set<String> functions) {
        this.functions = functions;
        this.treeRoot = context;
        this.functionName = treeRoot.getChild(1).getText().toUpperCase();
        this.classWriter = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
        //..
    }

    public byte[] compile() {
        classWriter.visit(V1_7, ACC_PUBLIC, functionName, null, "java/lang/Object", new String[]{"FunctionInterface"});
        classWriter.visitField(ACC_PRIVATE, "args", "[Ljava/lang/Object;", null, null).visitEnd();
        classWriter.visitField(ACC_PRIVATE, "count", "I", null, null).visitEnd();
        classWriter.visitField(ACC_PRIVATE, "value", "Ljava/lang/Object;", null, null).visitEnd();
        initVisitor = classWriter.visitMethod(ACC_PUBLIC, "<init>",
                "()V", null, null);
        initVisitor.visitCode();
        initVisitor.visitVarInsn(ALOAD, 0);
        initVisitor.visitMethodInsn(INVOKESPECIAL, "java/lang/Object", "<init>", "()V");
        initVisitor.visitVarInsn(ALOAD, 0);
        initVisitor.visitIntInsn(BIPUSH, 0);
        initVisitor.visitFieldInsn(PUTFIELD, functionName, "count", "I");

        //initVisitor.visitEnd();
        visit(treeRoot.lambda()); // visit lambda body
        classWriter.visitEnd();
        return classWriter.toByteArray();
    }

    @Override
    public Byte[] visitStaticLambdaObject(FunParser.StaticLambdaObjectContext ctx) {
        initVisitor.visitInsn(RETURN);
        initVisitor.visitMaxs(0, 0);
        initVisitor.visitEnd();

        classWriter.visitField(ACC_PRIVATE + ACC_STATIC, "staticValue", "Ljava/lang/Object;", null, null).visitEnd();
        MethodVisitor methodVisitor = classWriter.visitMethod(ACC_STATIC, "<clinit>", "()V", null, null);
        methodVisitor.visitCode();
        FunParser.ObjectContext object = ctx.object();
        String objectName = object.PATH().getText().substring(1);
        String[] args = object.SIGNATURE().getText().substring(1, object.SIGNATURE().getText().length() - 3).split("@");
        if (!args[0].equals("")) {
            unwrapArguments(args);
        }
        String signature = object.SIGNATURE().getText().replace("@", "");
        System.err.println("new Object: " + objectName);
        methodVisitor.visitTypeInsn(NEW, objectName);
        methodVisitor.visitInsn(DUP);
        if (!signature.substring(signature.length() - 2).equals(")V")) {
            visit(ctx.expression());
        }
        methodVisitor.visitMethodInsn(Opcodes.INVOKESPECIAL, objectName, "<init>", signature);
        methodVisitor.visitFieldInsn(PUTSTATIC, functionName, "staticValue", "Ljava/lang/Object;");
        methodVisitor.visitInsn(RETURN);
        methodVisitor.visitMaxs(0, 0);
        methodVisitor.visitEnd();
        getStaticValueCompiler();
        runStaticCompiler();
        return null;
    }

    private void runStaticCompiler() {
        methodVisitor = classWriter.visitMethod(ACC_PUBLIC, "run",
                "()LFunctionInterface;", null, null);
        methodVisitor.visitCode();
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitInsn(ARETURN);
        methodVisitor.visitMaxs(0, 0);
        methodVisitor.visitEnd();
        //To change body of created methods use File | Settings | File Templates.
    }

    private void getStaticValueCompiler() {
        methodVisitor = classWriter.visitMethod(ACC_PUBLIC, "getValue",
                "()Ljava/lang/Object;", null, null);
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitFieldInsn(GETSTATIC, this.functionName, "staticValue",
                Type.getDescriptor(Object.class));
        methodVisitor.visitInsn(ARETURN);
        methodVisitor.visitMaxs(0, 0);
        methodVisitor.visitEnd();

    }

    @Override
    public Byte[] visitLambdaWithoutArguments(FunParser.LambdaWithoutArgumentsContext ctx) {
        initVisitor.visitInsn(RETURN);
        initVisitor.visitMaxs(0, 0);
        initVisitor.visitEnd();
        this.argumentCount = 0;

        //add argument
        methodVisitor = classWriter.visitMethod(ACC_PUBLIC, "add",
                "(Ljava/lang/Object;)V", null, null);
        methodVisitor.visitCode();
        methodVisitor.visitInsn(RETURN);
        methodVisitor.visitMaxs(0, 0);
        methodVisitor.visitEnd();
        //get value
        getValueCompiler();
        //run
        runCompiler(ctx.expression());
        return null;
    }

    @Override
    public Byte[] visitLambdaWithArguments(FunParser.LambdaWithArgumentsContext ctx) {

        this.arguments = mapArgumentNames(ctx);
        this.argumentCount = getArgumentCount(ctx);
        if (argumentCount != arguments.size()) {
            System.err.println("Error: lambda has two arguments with same names");
            return null;
        }
        //create array
        initVisitor.visitVarInsn(ALOAD, 0);
        initVisitor.visitIntInsn(BIPUSH, argumentCount);
        initVisitor.visitTypeInsn(ANEWARRAY, "Ljava/lang/Object;");
        initVisitor.visitFieldInsn(PUTFIELD, functionName, "args", "[Ljava/lang/Object;");
        initVisitor.visitInsn(RETURN);
        initVisitor.visitMaxs(0, 0);
        initVisitor.visitEnd();
        //getValue method compile
        getValueCompiler();
        //addArg method compile
        addArgCompile();
        //run method compile
        runCompiler(ctx.getChild(ctx.getChildCount() - 1));

        return null;
    }

    private void runCompiler(ParseTree child) {
        methodVisitor = classWriter.visitMethod(ACC_PUBLIC, "run",
                "()LFunctionInterface;", null, null);
        methodVisitor.visitCode();

        Label run = new Label();
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitFieldInsn(GETFIELD, this.functionName, "count",
                Type.getDescriptor(int.class));
        methodVisitor.visitIntInsn(BIPUSH, argumentCount);
        methodVisitor.visitJumpInsn(IF_ICMPEQ, run);
        methodVisitor.visitVarInsn(ALOAD, 0);
        //methodVisitor.visitIntInsn(BIPUSH, 0);
        methodVisitor.visitInsn(ARETURN);
        methodVisitor.visitLabel(run);
        visit(child);
        methodVisitor.visitInsn(DUP);
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitInsn(SWAP);
        methodVisitor.visitFieldInsn(PUTFIELD, this.functionName, "value", "Ljava/lang/Object;");
        methodVisitor.visitInsn(ARETURN);
        methodVisitor.visitMaxs(0, 0);
        methodVisitor.visitEnd();
    }

    private void getValueCompiler() {
        methodVisitor = classWriter.visitMethod(ACC_PUBLIC, "getValue",
                "()Ljava/lang/Object;", null, null);
        methodVisitor.visitCode();
        Label returnResult = new Label();
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitFieldInsn(GETFIELD, this.functionName, "count",
                Type.getDescriptor(int.class));
        methodVisitor.visitIntInsn(BIPUSH, argumentCount);
        methodVisitor.visitJumpInsn(IF_ICMPEQ, returnResult);
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitInsn(ARETURN);
        methodVisitor.visitLabel(returnResult);
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitFieldInsn(GETFIELD, this.functionName, "value",
                Type.getDescriptor(Object.class));
        methodVisitor.visitInsn(ARETURN);
        methodVisitor.visitMaxs(0, 0);
        methodVisitor.visitEnd();
    }


    private void addArgCompile() {
        methodVisitor = classWriter.visitMethod(ACC_PUBLIC, "add", "(Ljava/lang/Object;)V", null, null);
        methodVisitor.visitCode();
        Label addArg = new Label();
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitFieldInsn(GETFIELD, this.functionName, "count",
                Type.getDescriptor(int.class));
        methodVisitor.visitIntInsn(BIPUSH, argumentCount);
        methodVisitor.visitJumpInsn(IF_ICMPLT, addArg);
        methodVisitor.visitInsn(RETURN);
        //return false
        methodVisitor.visitLabel(addArg);
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitFieldInsn(GETFIELD, this.functionName, "args", "[Ljava/lang/Object;");
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitFieldInsn(GETFIELD, this.functionName, "count", "I");
        methodVisitor.visitIntInsn(ALOAD, 1);
        methodVisitor.visitInsn(AASTORE);

        //++count
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitFieldInsn(GETFIELD, this.functionName, "count", "I");
        methodVisitor.visitIntInsn(BIPUSH, 1);
        methodVisitor.visitInsn(IADD);
        methodVisitor.visitFieldInsn(PUTFIELD, this.functionName, "count", "I");
        methodVisitor.visitInsn(RETURN);

        methodVisitor.visitMaxs(0, 0);
        methodVisitor.visitEnd();
    }


    private void eval(ParserRuleContext ctx) {
        visit(ctx.getChild(0));
        visit(ctx.getChild(2));
    }

    @Override
    public Byte[] visitAddSub(FunParser.AddSubContext ctx) {
        eval(ctx);
        int operator = ctx.operator.getType();
        System.err.println("add/sub");
        switch (operator) {
            case FunLexer.PLUS:
                convertTwoNumbersToInt(methodVisitor);
                methodVisitor.visitInsn(IADD);
                saveIntToNumber(methodVisitor);
                break;

            case FunLexer.MINUS:
                convertTwoNumbersToInt(methodVisitor);
                methodVisitor.visitInsn(ISUB);
                saveIntToNumber(methodVisitor);
                break;
        }
        return null;
    }

    private void saveIntToNumber(MethodVisitor methodVisitor) {
        methodVisitor.visitTypeInsn(NEW, "Number");
        methodVisitor.visitInsn(DUP);
        methodVisitor.visitVarInsn(ASTORE, 1);
        methodVisitor.visitInsn(SWAP);
        methodVisitor.visitMethodInsn(INVOKESPECIAL, "Number", "<init>", "(I)V");
        methodVisitor.visitVarInsn(ALOAD, 1);
    }

    @Override
    public Byte[] visitNewClass(FunParser.NewClassContext ctx) {
        FunParser.ObjectContext object = ctx.object();
        String objectName = object.PATH().getText().substring(1);
        String[] args = object.SIGNATURE().getText().substring(1, object.SIGNATURE().getText().length() - 3).split("@");
        if (!args[0].equals("")) {
            unwrapArguments(args);
        }

        String signature = object.SIGNATURE().getText().replace("@", "");
        System.err.println("new Object: " + objectName);
        methodVisitor.visitTypeInsn(NEW, "ObjectWrapper");
        methodVisitor.visitInsn(DUP);
        methodVisitor.visitTypeInsn(NEW, objectName);
        methodVisitor.visitInsn(DUP);
        methodVisitor.visitMethodInsn(Opcodes.INVOKESPECIAL, objectName, "<init>", signature);
        methodVisitor.visitMethodInsn(Opcodes.INVOKESPECIAL, "ObjectWrapper", "<init>", "(Ljava/lang/Object;)V");
        return null;
    }


    @Override
    public Byte[] visitMethodCall(FunParser.MethodCallContext ctx) {
        methodVisitor.visitTypeInsn(NEW, "ObjectWrapper");
        methodVisitor.visitInsn(DUP);
        visit(ctx.expression());
        String objectName = ctx.PATH().getText().substring(1);
        String funSignature = ctx.object().SIGNATURE().getText();

        String[] args = funSignature.split("@");
        args[args.length - 1] = "";


        String functionName = ctx.object().PATH().getText().substring(1);
        String signature = ctx.object().SIGNATURE().getText().replace("@", "");
        System.err.print("Method call: " + objectName + "." + functionName + " " + signature);
        methodVisitor.visitMethodInsn(Opcodes.INVOKEINTERFACE, "FunctionInterface", "getValue",
                "()Ljava/lang/Object;");
        methodVisitor.visitTypeInsn(CHECKCAST, objectName);
        methodVisitor.visitMethodInsn(Opcodes.INVOKEVIRTUAL, objectName, functionName, signature);
        String returnType = signature.substring(signature.length() - 2);
        if (returnType.equals(")V")) {
            methodVisitor.visitMethodInsn(Opcodes.INVOKESPECIAL, "ObjectWrapper", "<init>", "()V");
        } else {
            methodVisitor.visitMethodInsn(Opcodes.INVOKESPECIAL, "ObjectWrapper", "<init>", "(Ljava/lang/Object;)V");
        }

        return null;
    }

    /*@Override
    public Byte[] visitObjectRemember(FunParser.ObjectRememberContext ctx) {

        FunParser.ObjectContext object = ctx.object();
        String objectName = object.PATH().getText().substring(1);
        String[] args = object.SIGNATURE().getText().substring(1, object.SIGNATURE().getText().length() - 3).split("@");
        if (!args[0].equals("")) {
            unwrapArguments(args);
        }

        String signature = object.SIGNATURE().getText().replace("@", "");
        classWriter.visitField(ACC_PRIVATE, objectName, "ObjectWrapper", null, null).visitEnd();
        System.err.println("remember Object: " + objectName);
        methodVisitor.visitTypeInsn(NEW, "ObjectWrapper");
        methodVisitor.visitInsn(DUP);
        methodVisitor.visitTypeInsn(NEW, objectName);
        methodVisitor.visitInsn(DUP);
        methodVisitor.visitMethodInsn(Opcodes.INVOKESPECIAL, objectName, "<init>", signature);
        methodVisitor.visitMethodInsn(Opcodes.INVOKESPECIAL, "ObjectWrapper", "<init>", "(Ljava/lang/Object;)V");
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitInsn(SWAP);
        methodVisitor.visitFieldInsn(PUTFIELD, functionName, objectName, "ObjectWrapper");
        return null;
    }    */

    private void unwrapArguments(String[] args) {
        for (int i = args.length - 1; i >= 0; --i) {
            String signature = args[i];
            if (signature.equals(""))
                break;
            if (signature.charAt(0) == 'L') {
                signature = signature.substring(1, signature.length() - 1);
            }
            methodVisitor.visitTypeInsn(CHECKCAST, signature);
            methodVisitor.visitVarInsn(ASTORE, i + 1);
        }
        for (int i = 0; i < args.length - 1; ++i) {
            methodVisitor.visitVarInsn(ALOAD, i + 1);
        }
    }


    @Override
    public Byte[] visitMulDiv(FunParser.MulDivContext ctx) {
        eval(ctx);
        System.err.println("mul/div");
        int operator = ctx.operator.getType();
        switch (operator) {
            case FunLexer.MULT:
                convertTwoNumbersToInt(methodVisitor);
                methodVisitor.visitInsn(IMUL);
                saveIntToNumber(methodVisitor);
                break;

            case FunLexer.DIV:
                convertTwoNumbersToInt(methodVisitor);
                methodVisitor.visitInsn(IDIV);
                saveIntToNumber(methodVisitor);
                break;

            case FunLexer.MOD:
                convertTwoNumbersToInt(methodVisitor);
                methodVisitor.visitInsn(IREM);
                saveIntToNumber(methodVisitor);
                break;
        }

        return null;
    }

    private void convertTwoNumbersToInt(MethodVisitor methodVisitor) {
        methodVisitor.visitInsn(SWAP);
        methodVisitor.visitMethodInsn(Opcodes.INVOKEINTERFACE, "FunctionInterface", "getValue",
                "()Ljava/lang/Object;");
        methodVisitor.visitTypeInsn(CHECKCAST, "java/lang/Integer");
        methodVisitor.visitMethodInsn(INVOKEVIRTUAL, "java/lang/Integer", "intValue", "()I");
        methodVisitor.visitInsn(SWAP);
        methodVisitor.visitMethodInsn(Opcodes.INVOKEINTERFACE, "FunctionInterface", "getValue",
                "()Ljava/lang/Object;");
        methodVisitor.visitTypeInsn(CHECKCAST, "java/lang/Integer");
        methodVisitor.visitMethodInsn(INVOKEVIRTUAL, "java/lang/Integer", "intValue", "()I");
    }

    @Override
    public Byte[] visitIfThenElse(FunParser.IfThenElseContext ctx) {
        visit(ctx.expression(0));
        convertNumbersToInt(methodVisitor);
        Label label = new Label();
        methodVisitor.visitJumpInsn(IFEQ, label);   //if eval of expression "returns" 0,    than expression==false
        visit(ctx.expression(1));
        Label end = new Label();
        methodVisitor.visitJumpInsn(GOTO, end);
        methodVisitor.visitLabel(label);
        visit(ctx.expression(2));
        methodVisitor.visitLabel(end);
        return null;
    }

    private void convertNumbersToInt(MethodVisitor methodVisitor) {
        methodVisitor.visitMethodInsn(Opcodes.INVOKEINTERFACE, "FunctionInterface", "getValue",
                "()Ljava/lang/Object;");
        methodVisitor.visitTypeInsn(CHECKCAST, "java/lang/Integer");
        methodVisitor.visitMethodInsn(INVOKEVIRTUAL, "java/lang/Integer", "intValue", "()I");
    }


    @Override
    public Byte[] visitLogic(FunParser.LogicContext ctx) {
        eval(ctx);
        int operator = ctx.operator.getType();
        convertTwoNumbersToInt(methodVisitor);

        switch (operator) {
            case FunParser.AND: {
                Label returnFalse1 = new Label();
                methodVisitor.visitJumpInsn(IFEQ, returnFalse1);
                Label returnFalse2 = new Label();
                methodVisitor.visitJumpInsn(IFEQ, returnFalse2);
                insertNumber(methodVisitor, 1);
                Label end = new Label();
                methodVisitor.visitJumpInsn(GOTO, end);
                methodVisitor.visitLabel(returnFalse1);
                methodVisitor.visitInsn(POP);
                methodVisitor.visitLabel(returnFalse2);
                insertNumber(methodVisitor, 0);
                methodVisitor.visitLabel(end);
                break;
            }

            case FunParser.OR: {
                Label returnTrue1 = new Label();
                methodVisitor.visitJumpInsn(IFNE, returnTrue1);
                Label returnTrue2 = new Label();
                methodVisitor.visitJumpInsn(IFNE, returnTrue2);
                insertNumber(methodVisitor, 0);
                Label end = new Label();
                methodVisitor.visitJumpInsn(GOTO, end);
                methodVisitor.visitLabel(returnTrue1);
                methodVisitor.visitInsn(POP);
                methodVisitor.visitLabel(returnTrue2);
                insertNumber(methodVisitor, 1);
                methodVisitor.visitLabel(end);
                break;
            }

        }
        return null;
    }

    private void insertNumber(MethodVisitor methodVisitor, int i) {
        methodVisitor.visitTypeInsn(NEW, "Number");
        methodVisitor.visitInsn(DUP);
        methodVisitor.visitLdcInsn(i);
        methodVisitor.visitMethodInsn(INVOKESPECIAL, "Number", "<init>", "(I)V");
    }

    @Override
    public Byte[] visitCompare(FunParser.CompareContext ctx) {
        eval(ctx);
        convertTwoNumbersToInt(methodVisitor);
        int operator = ctx.operator.getType();
        int opcode;
        switch (operator) {
            case FunParser.EQUAL:
                opcode = IF_ICMPNE;
                break;

            case FunParser.NOT_EQUAL:
                opcode = IF_ICMPEQ;
                break;

            case FunParser.LESS:
                opcode = IF_ICMPGE;
                break;

            case FunParser.LESS_OR_EQUAL:
                opcode = IF_ICMPGT;
                break;

            case FunParser.MORE_OR_EQUAL:
                opcode = IF_ICMPLT;
                break;

            case FunParser.MORE:
                opcode = IF_ICMPLE;
                break;

            default:
                throw new NotImplementedException();
        }

        Label returnFalse = new Label();
        methodVisitor.visitJumpInsn(opcode, returnFalse);
        insertNumber(methodVisitor, 1);
        Label end = new Label();
        methodVisitor.visitJumpInsn(GOTO, end);
        methodVisitor.visitLabel(returnFalse);
        insertNumber(methodVisitor, 0);
        methodVisitor.visitLabel(end);
        return null;
    }

    @Override
    public Byte[] visitIdentifier(FunParser.IdentifierContext ctx) {
        String func = ctx.getText();
        if (arguments != null && arguments.containsKey(func)) {
            Integer argNumber = arguments.get(func);
            methodVisitor.visitVarInsn(ALOAD, 0);
            methodVisitor.visitFieldInsn(GETFIELD, this.functionName, "args", "[Ljava/lang/Object;");
            methodVisitor.visitLdcInsn(argNumber);
            methodVisitor.visitInsn(AALOAD);
            return null;
        }
        func = func.toUpperCase();

        if (functions.contains(func)) {
            methodVisitor.visitTypeInsn(NEW, func);
            methodVisitor.visitInsn(DUP);
            methodVisitor.visitMethodInsn(Opcodes.INVOKESPECIAL, func, "<init>", "()V");
            methodVisitor.visitMethodInsn(INVOKEINTERFACE, "FunctionInterface", "run", "()LFunctionInterface;");
            return null;
        }

        return super.visitIdentifier(ctx);
    }

    @Override
    public Byte[] visitApplications(FunParser.ApplicationsContext ctx) {
        FunParser.ApplicationsListContext app = ctx.applicationsList();
        //visit(app);
        if (app.application.unary) {
            visit(app.application.func);
            return null;
        }
        ApplicationTree cur = app.application;
        apply(cur);
        methodVisitor.visitMethodInsn(INVOKEINTERFACE, "FunctionInterface", "run", "()LFunctionInterface;");
        return null;
    }

    public void apply(ApplicationTree application) {
        if (application.unary) {
            visit(application.func);
            return;
        }
        apply(application.next);
        methodVisitor.visitInsn(DUP);
        visit(application.func);
        methodVisitor.visitMethodInsn(INVOKEINTERFACE, "FunctionInterface", "add", "(Ljava/lang/Object;)V");
    }


    @Override
    public Byte[] visitLiteral(FunParser.LiteralContext ctx) {
        System.err.println("bipush");
        insertNumber(methodVisitor, new Integer(ctx.LITERAL().toString()));
        return null;
    }

    public String getName() {
        return this.functionName;
    }

    private String getArgumentNameByIndex(Integer index) {
        return "arg" + index.toString();
    }

    private Map<String, Integer> mapArgumentNames(FunParser.LambdaWithArgumentsContext ctx) {
        Map<String, Integer> result = new TreeMap<>();
        int count = 0;
        for (TerminalNode node : ctx.IDENTIFIER()) {
            result.put(node.getText(), count);
            ++count;
        }
        return result;
    }

    int getArgumentCount(FunParser.LambdaWithArgumentsContext ctx) {
        return ctx.IDENTIFIER().size();
    }


    private FunParser.LetContext treeRoot;
    private String functionName;
    private Integer argumentCount = 0;
    MethodVisitor methodVisitor;
    MethodVisitor initVisitor;
    ClassWriter classWriter;
    Map<String, Integer> arguments;

}

